import unittest

import cv2
import numpy as np
from ocr_process.infer.infer_engine import OrtInferSession

from taskinfo.table_process import bbox_decode, decode_by_ind, nms, bbox_post_process, sigmoid
from taskinfo.uv_pipeline import UvPipeline


class UVTestPipeline(unittest.TestCase):
    def test_uv_pipeline(self):
        uv_predict = UvPipeline(model_path=r"D:\ocr\onnx\deploy\uvdoc\inference.onnx", use_openvino=False)
        img = cv2.imread(r"D:\ocr\train_data\multi_class\multi_0520\rong_cheng\Y00092410173565_1.jpg")
        result = uv_predict(img)
        if result['output_imgs']:
            image_data = result['output_imgs'][0]
            cv2.imwrite(r"test_result.jpg", image_data)
